
from ultralytics import YOLO


def main(args):
    kwargs = {
    }

    model = YOLO(args.model)

    results = model.train(
        **kwargs,
        # 结果位于 project/name 目录中
        project='runs/detect',
        name='dotasplit_train',
        # 是否覆盖
        exist_ok=True,

        # 训练参数
        #

        # 数据集配置文件
        data=args.data,
        # 数据集用于训练的比例
        fraction=1.0,
        # true to 使用预训练的权重
        pretrained=False,
        # 迭代
        epochs=10,
        # 连续多少轮没有改善则提前终止
        patience=5,
        # true to 保存检查点
        save=True,
        # 多少轮保存一次
        save_period=2,
        # true to 使用保存点，这时 args.model 应该指向 weights\last.pt
        resume=args.resume,
        # 随机种子
        seed=0,
        # 启用确定性
        deterministic=True,
        # 每批图像数量
        batch=8,
        # 训练使用 640x640
        imgsz=640,
        # true to 接受矩形输入，这里不需要
        rect=False,
        # 数据加载线程
        workers=20,
        # true to cache using ram, disk
        cache=True,
        # true to 使用 FP16 训练
        half=False,

        # true to 启用性能分析
        profile=True,

        # 评估参数
        #

        # true to 每轮训练后评估
        val=False,
        # 使用何种数据集做评估
        split='val',
        save_json=True,

        device=args.device,
        verbose=args.verbose)
    print(results)

    # success = model.export(format='onnx')
    # print(success)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('model', type=str)
    parser.add_argument('data', type=str)
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--device', type=str, default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--verbose', action='store_true')
    args = parser.parse_args()
    main(args)
